
"""
将KITTI的标签文件（000001.txt ...）转换成json格式（COCO format）。
"""

import os
import cv2
import json
from collections import OrderedDict

images, categories, annotations = [], [], []

category_dict = {
    "Car": 0,
    "Pedestrian": 1,
    "Cyclist": 2
}

# categories info
"""
{
    "id": int,
    "name": str,
    "supercategory": str,
}
"""
for cat_n in category_dict:
    cat_info = OrderedDict()
    cat_info["supercategory"] = ""
    cat_info["id"] = category_dict[cat_n]
    cat_info["name"] = cat_n
    # categories.append({"supercategory": "", "id": category_dict[cat_n], "name": cat_n})
    categories.append(cat_info)

# print(categories)

txt_root_path = 'E:\\0_My_codes\\ws_subl\\0217_yolov3_results\\label_2_ori_eval\\'
img_root_path = 'E:\\0_My_codes\\ws_pycharm\\non_syn_projs\\kitti_2d_proc\\kittivoc_ori_1_1\\kittivoc_test\\JPEGImages\\'

txt_name_list = []

for f_name in os.listdir(txt_root_path):
    # print(os.path.splitext(f_name))
    if os.path.splitext(f_name)[1] == '.txt':
        f_name = os.path.splitext(f_name)[0]
        txt_name_list.append(f_name)

anno_id = 0
for item in txt_name_list:
    # print(item)

    # img
    img = cv2.imread(img_root_path + item + '.jpg')
    [height, width, _] = img.shape
    # images info
    """
    image{
        "id": int,
        "width": int,
        "height": int,
        "file_name": str,
        "license": int,
        "flickr_url": str,
        "coco_url": str,
        "date_captured": datetime,
    }
    """
    img_name = item + '.jpg'
    img_id = int(item)
    img_info = OrderedDict()
    img_info["id"] = img_id
    img_info["width"] = width
    img_info["height"] = height
    img_info["file_name"] = img_name
    # images.append({"file_name": img_name, "height": height, "width": width, "id": img_id})
    images.append(img_info)

    # annotation info
    """
    annotation{
        "id": int,
        "image_id": int,
        "category_id": int,
        "segmentation": RLE or [polygon],
        "area": float,
        "bbox": [x,y,width,height],
        "iscrowd": 0 or 1,
    }
    """

    with open(txt_root_path + item + '.txt', 'r') as f:
        for line in f.readlines():
            # print(line)
            line = line.strip().split(' ')
            # print(line)
            # transform---------------
            img_id = int(item)
            cat_id = category_dict[line[0]]
            x = float(line[4])
            y = float(line[5])
            w = float("%.2f" % (float(line[6]) - float(line[4])))
            h = float("%.2f" % (float(line[7]) - float(line[5])))
            area = float("%.2f" % (w * h))
            bbox = [x, y, w, h]
            # score = float(line[])
            # transform---------------

            anno_info = OrderedDict()
            anno_info["id"] = anno_id
            anno_info["image_id"] = img_id
            anno_info["category_id"] = cat_id
            anno_info["area"] = area
            anno_info["bbox"] = bbox
            anno_info["iscrowd"] = 0
            # anno_info = {'id': anno_id, 'image_id': img_id, 'category_id': cat_id, 'area': area, 'bbox': bbox, 'iscrowd': 0}
            annotations.append(anno_info)
            anno_id += 1

all_json = {"images": images, "annotations": annotations, "categories": categories}

with open("gt.json", "w") as outfile:
    json.dump(all_json, outfile)

print('done.')